from types import SimpleNamespace

from wenet.zipformer.decoder_attention import TransformerDecoder, AttentionDecoderModel

def init_decoder4zipformer(config_zipformer):
    params = SimpleNamespace(**config_zipformer)
    decoder = AttentionDecoderModel(
        vocab_size=params.vocab_size,
        decoder_dim=params.attention_decoder_dim,
        num_decoder_layers=params.attention_decoder_num_layers,
        attention_dim=params.attention_decoder_attention_dim,
        num_heads=int(params.attention_decoder_num_heads),
        feedforward_dim=params.attention_decoder_feedforward_dim,
        memory_dim=max(params.encoder_dim),
        sos_id=params.sos_id,
        eos_id=params.eos_id,
        ignore_id=params.ignore_id,
        label_smoothing=params.label_smoothing,
    )
    return decoder
    # param_obj = SimpleNamespace(**config_zipformer)
    # decoder = TransformerDecoder(
    #     vocab_size=param_obj.vocab_size,
    #     d_model=param_obj.decoder_dim,
    #     num_decoder_layers=param_obj.attention_decoder_num_layers,
    #     attention_dim=param_obj.attention_decoder_attention_dim,
    #     num_heads=param_obj.attention_decoder_num_heads,
    #     feedforward_dim=param_obj.attention_decoder_feedforward_dim,
    #     memory_dim=max(param_obj.encoder_dim),
    #     dropout=0.1,
    # )
    # return decoder